"""Import Packages"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchview import draw_graph # Ref and credits: https://github.com/mert-kurttutan/torchview for visualizing model graph.
from torchvision.utils import make_grid
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as T
from torchsummary import summary
from PIL import Image
from DDPM_Scheduler import DDPM_Scheduler
from Unet_module import UNET
import yaml
import cv2
%matplotlib inline

"""Extraction of Parameters from YAML file."""

with open('config.yaml', 'r') as file:
    try:
        config = yaml.safe_load(file)
    except yaml.YAMLError as exc:
        print(exc)
    print(config)

{'version': 'v1.0', 'image_size': {'channels': 3, 'height': 32, 'width': 32}, 'sample_dir': 'Generated', 'seed': 34, 'diffusion_process': {'num_training_steps': 1000, 'num_inference_steps': 1000, 'beta_1': 0.0001, 'beta_T': 0.02}, 'reverse_process': {'time_emb_dim': 128, 'lr': 0.001}}

Diffusion_config = config["diffusion_process"]
Model_config = config["reverse_process"]
Image_dim = config["image_size"]
seed = config["seed"]
sample_dir = config["sample_dir"]

batch_size = 64
image_size = (Image_dim["channels"], Image_dim["height"], Image_dim["width"])
num_training_steps = Diffusion_config["num_training_steps"]
num_inference_steps = Diffusion_config["num_inference_steps"]
time_emb_dim = Model_config["time_emb_dim"]
lr = Model_config["lr"]

"""Images Directory"""
# Dataset Ref: https://www.kaggle.com/datasets/mostafamozafari/bitmoji-faces

Root_Directory = 'D:\Envs\DCGAN\Bitmoji_Faces\BitmojiDataset'
for l in os.listdir(Root_Directory):
    if os.path.isdir(os.path.join(Root_Directory, ''.join(l))):
        Directory = os.path.join(Root_Directory, ''.join(l))
Bitmoji_Images = os.listdir(Directory)
len(Bitmoji_Images)

4084

"""Class CSV with Image_ID and Gender Class."""

csv_directory = os.path.join(Root_Directory, 'attributes.csv')
Class = pd.read_csv(csv_directory)

"""Rename Target Column."""

print(Class.head())
Class = Class.rename(columns = {'is_male': 'Gender'})
Class.head()

   image_id  is_male
0  0000.jpg        1
1  0001.jpg       -1
2  0002.jpg        1
3  0003.jpg       -1
4  0004.jpg       -1

"""Replace -1 to 0 for Female Category."""

Class['Gender'] = Class['Gender'].apply(lambda x: 0 if x == -1 else x) # 1 - Male, 0 - Female
Class.head()

"""Creating Custom Dataset with Images and Labels from CSV file."""

class Image_Label_Dataset(Dataset):
    def __init__(self, Image_folder, Class_csv, transforms):
        self.df = Class_csv
        self.images_folder = Image_folder
        self.transform = transforms

    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        filename = self.df['image_id'][idx]
        label = self.df['Gender'][idx]
        image = Image.open(os.path.join(self.images_folder, filename))
        image = self.transform(image)
        image = (image * 2) - 1
        return image, label

"""Transformations for Image."""

transform = T.Compose([
                      T.Resize((image_size[1], image_size[2])), 
                      T.ToTensor()
                      ])

Dataset = Image_Label_Dataset(Directory, Class, transform)

"""Functions to show and save images in batch."""

def show_and_save_images(images, nmax = 64, t: int = None, index = None, save = False, show = True):
    
    images = torch.clamp(images, -1, 1).detach().cpu()
    images = (images + 1) / 2
    images = (images * 255).type(torch.uint8)
    fig, ax = plt.subplots(figsize = (8, 8))
    ax.set_xticks([]); ax.set_yticks([])
    ax.imshow(make_grid(images[:nmax], nrow = 8).permute(1, 2, 0))
    if t is not None:
        ax.set_title('Noise at Timestep: {}'.format(t))
    
    if show:  
        plt.show()
    if save:
        img_name = f'generated-image-step-{index:0=4d}-{t:0=3d}.png'
        fig.savefig(os.path.join(sample_dir, img_name), bbox_inches='tight')
        plt.close()
        

def show_batch(dl, nmax = 64):
    for images, _ in dl:
        show_and_save_images(images, nmax)
        break

"""Creating DataLoader from above dataset."""

Train_Dl = DataLoader(Dataset, batch_size, shuffle = True, num_workers = 0)

Label = {"female" : 0, "male" : 1}
Label_size = len(Label)
print(Label_size)

2

img, _ = next(iter(Train_Dl))

img.size()

torch.Size([64, 3, 32, 32])

show_batch(Train_Dl)

"""Initialising Scheduler."""

ddpm = DDPM_Scheduler(Diffusion_config["beta_1"], Diffusion_config["beta_T"], num_training_steps)

"""Adding Noise to the images for the given timesteps."""

t = torch.Tensor([0, 50, 100, 150, 200, 300, 600, 700, 800, 900, 999]).long()
Images = []
for step in t:
    Img_t = []
    for image in img:
        noise = torch.rand_like((image))
        noised_image = ddpm.diffusion_process(image.unsqueeze(0), step, noise.unsqueeze(0))
        Img_t.append(noised_image)
    Images.append(torch.cat(Img_t))

Images[0].size()

torch.Size([64, 3, 32, 32])

"""Visualization of Images with Noise at each timestep."""

for i, step in enumerate(t):
    show_and_save_images(Images[i], nmax = 64, t = int(step), index = i, save = True)
    print()

def get_default_device():
    """Use GPU if available, else CPU"""
    if torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            print(torch.cuda.get_device_properties(i))
        return torch.device('cuda')
    else:
        return torch.device('cpu')
    
def to_device(data, device):
    """Move Tensors to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)

class DeviceDataLoader():
    """Wrap a Dataloader to move data to a device"""
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
        
    def __iter__(self):
        """Yield a batch of data after moving it to device"""
        for b in self.dl: 
            yield to_device(b, self.device)

    def __len__(self):
        """Number of batches"""
        return len(self.dl)

"""Check for the device."""
device = get_default_device()

_CudaDeviceProperties(name='NVIDIA GeForce RTX 3060 Laptop GPU', major=8, minor=6, total_memory=6143MB, multi_processor_count=30)

generator = torch.Generator(device = device)
generator.manual_seed(seed)

<torch._C.Generator at 0x20a8aa75e50>

"""Changing the DataLoader to the Device."""
Train_Dl = DeviceDataLoader(Train_Dl, device)

"""Initialising of UNET Class."""

unet = UNET(image_size[0], time_emb_dim, Label_size)

Num_of_parameters = sum(p.numel() for p in unet.parameters())
print("Model Parameters : {:.3f} M".format(Num_of_parameters / 1e6))

Model Parameters : 1.589 M

img_dummy = torch.randn(3, 3, 32, 32)
labels_dummy = torch.randint(0, 2, size = (3,), dtype = torch.long)
t = torch.randint(0, num_training_steps, (3,)).type(torch.long)

"""Summary report of Initialised UNET model."""
summary(unet, [(img_dummy), (labels_dummy), (t)], device = 'cpu')

===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
├─Time_Encoding: 1-1                          [-1, 128]                 --
├─Embedding: 1-2                              [-1, 128]                 256
├─Sequential: 1-3                             [-1, 128]                 --
|    └─Linear: 2-1                            [-1, 256]                 33,024
|    └─SiLU: 2-2                              [-1, 256]                 --
|    └─Linear: 2-3                            [-1, 128]                 32,896
├─Conv2d: 1-4                                 [-1, 16, 32, 32]          448
├─ModuleList: 1                               []                        --
|    └─Resnet_Block: 2-4                      [-1, 32, 32, 32]          --
|    |    └─Sequential: 3-1                   [-1, 32]                  4,128
|    |    └─Conv_block: 3-2                   [-1, 32, 32, 32]          4,704
|    |    └─Conv_block: 3-3                   [-1, 32, 32, 32]          9,312
|    |    └─Conv2d: 3-4                       [-1, 32, 32, 32]          544
|    └─UNET_Attention: 2-5                    [-1, 32, 32, 32]          --
|    |    └─LayerNorm: 3-5                    [-1, 1024, 32]            64
|    |    └─MultiHeadAttention: 3-6           [-1, 1024, 32]            4,096
|    |    └─LayerNorm: 3-7                    [-1, 1024, 32]            64
|    |    └─Linear: 3-8                       [-1, 1024, 64]            2,048
|    └─Downsample: 2-6                        [-1, 32, 16, 16]          --
|    |    └─Conv2d: 3-9                       [-1, 32, 16, 16]          9,248
|    └─Resnet_Block: 2-7                      [-1, 64, 16, 16]          --
|    |    └─Sequential: 3-10                  [-1, 64]                  8,256
|    |    └─Conv_block: 3-11                  [-1, 64, 16, 16]          18,624
|    |    └─Conv_block: 3-12                  [-1, 64, 16, 16]          37,056
|    |    └─Conv2d: 3-13                      [-1, 64, 16, 16]          2,112
|    └─UNET_Attention: 2-8                    [-1, 64, 16, 16]          --
|    |    └─LayerNorm: 3-14                   [-1, 256, 64]             128
|    |    └─MultiHeadAttention: 3-15          [-1, 256, 64]             16,384
|    |    └─LayerNorm: 3-16                   [-1, 256, 64]             128
|    |    └─Linear: 3-17                      [-1, 256, 128]            8,192
|    └─Downsample: 2-9                        [-1, 64, 8, 8]            --
|    |    └─Conv2d: 3-18                      [-1, 64, 8, 8]            36,928
|    └─Resnet_Block: 2-10                     [-1, 128, 8, 8]           --
|    |    └─Sequential: 3-19                  [-1, 128]                 16,512
|    |    └─Conv_block: 3-20                  [-1, 128, 8, 8]           74,112
|    |    └─Conv_block: 3-21                  [-1, 128, 8, 8]           147,840
|    |    └─Conv2d: 3-22                      [-1, 128, 8, 8]           8,320
|    └─UNET_Attention: 2-11                   [-1, 128, 8, 8]           --
|    |    └─LayerNorm: 3-23                   [-1, 64, 128]             256
|    |    └─MultiHeadAttention: 3-24          [-1, 64, 128]             65,536
|    |    └─LayerNorm: 3-25                   [-1, 64, 128]             256
|    |    └─Linear: 3-26                      [-1, 64, 256]             32,768
|    └─Conv2d: 2-12                           [-1, 128, 8, 8]           147,584
├─ModuleList: 1                               []                        --
|    └─Resnet_Block: 2-13                     [-1, 128, 8, 8]           --
|    |    └─Sequential: 3-27                  [-1, 128]                 16,512
|    |    └─Conv_block: 3-28                  [-1, 128, 8, 8]           147,840
|    |    └─Conv_block: 3-29                  [-1, 128, 8, 8]           147,840
|    |    └─Identity: 3-30                    [-1, 128, 8, 8]           --
|    └─UNET_Attention: 2-14                   [-1, 128, 8, 8]           --
|    |    └─LayerNorm: 3-31                   [-1, 64, 128]             256
|    |    └─MultiHeadAttention: 3-32          [-1, 64, 128]             65,536
|    |    └─LayerNorm: 3-33                   [-1, 64, 128]             256
|    |    └─Linear: 3-34                      [-1, 64, 256]             32,768
|    └─Resnet_Block: 2-15                     [-1, 64, 8, 8]            --
|    |    └─Sequential: 3-35                  [-1, 64]                  8,256
|    |    └─Conv_block: 3-36                  [-1, 64, 8, 8]            73,920
|    |    └─Conv_block: 3-37                  [-1, 64, 8, 8]            37,056
|    |    └─Conv2d: 3-38                      [-1, 64, 8, 8]            8,256
├─ModuleList: 1                               []                        --
|    └─ModuleList: 2                          []                        --
|    |    └─Resnet_Block: 3-39                [-1, 64, 8, 8]            168,448
|    |    └─UNET_Attention: 3-40              [-1, 64, 8, 8]            24,832
|    |    └─ConvTranspose2d: 3-41             [-1, 64, 8, 8]            36,928
|    └─ModuleList: 2                          []                        --
|    |    └─Resnet_Block: 3-42                [-1, 32, 8, 8]            54,528
|    |    └─UNET_Attention: 3-43              [-1, 32, 8, 8]            6,272
|    |    └─Upsample: 3-44                    [-1, 32, 16, 16]          16,416
|    └─ModuleList: 2                          []                        --
|    |    └─Resnet_Block: 3-45                [-1, 16, 16, 16]          14,720
|    |    └─UNET_Attention: 3-46              [-1, 16, 16, 16]          1,600
|    |    └─Upsample: 3-47                    [-1, 16, 32, 32]          4,112
├─GroupNorm: 1-5                              [-1, 16, 32, 32]          32
├─SiLU: 1-6                                   [-1, 16, 32, 32]          --
├─Conv2d: 1-7                                 [-1, 3, 32, 32]           435
===============================================================================================
Total params: 1,588,643
Trainable params: 1,588,643
Non-trainable params: 0
Total mult-adds (M): 32.79
===============================================================================================
Input size (MB): 0.04
Forward/backward pass size (MB): 5.51
Params size (MB): 6.06
Estimated Total Size (MB): 11.61
===============================================================================================

===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
├─Time_Encoding: 1-1                          [-1, 128]                 --
├─Embedding: 1-2                              [-1, 128]                 256
├─Sequential: 1-3                             [-1, 128]                 --
|    └─Linear: 2-1                            [-1, 256]                 33,024
|    └─SiLU: 2-2                              [-1, 256]                 --
|    └─Linear: 2-3                            [-1, 128]                 32,896
├─Conv2d: 1-4                                 [-1, 16, 32, 32]          448
├─ModuleList: 1                               []                        --
|    └─Resnet_Block: 2-4                      [-1, 32, 32, 32]          --
|    |    └─Sequential: 3-1                   [-1, 32]                  4,128
|    |    └─Conv_block: 3-2                   [-1, 32, 32, 32]          4,704
|    |    └─Conv_block: 3-3                   [-1, 32, 32, 32]          9,312
|    |    └─Conv2d: 3-4                       [-1, 32, 32, 32]          544
|    └─UNET_Attention: 2-5                    [-1, 32, 32, 32]          --
|    |    └─LayerNorm: 3-5                    [-1, 1024, 32]            64
|    |    └─MultiHeadAttention: 3-6           [-1, 1024, 32]            4,096
|    |    └─LayerNorm: 3-7                    [-1, 1024, 32]            64
|    |    └─Linear: 3-8                       [-1, 1024, 64]            2,048
|    └─Downsample: 2-6                        [-1, 32, 16, 16]          --
|    |    └─Conv2d: 3-9                       [-1, 32, 16, 16]          9,248
|    └─Resnet_Block: 2-7                      [-1, 64, 16, 16]          --
|    |    └─Sequential: 3-10                  [-1, 64]                  8,256
|    |    └─Conv_block: 3-11                  [-1, 64, 16, 16]          18,624
|    |    └─Conv_block: 3-12                  [-1, 64, 16, 16]          37,056
|    |    └─Conv2d: 3-13                      [-1, 64, 16, 16]          2,112
|    └─UNET_Attention: 2-8                    [-1, 64, 16, 16]          --
|    |    └─LayerNorm: 3-14                   [-1, 256, 64]             128
|    |    └─MultiHeadAttention: 3-15          [-1, 256, 64]             16,384
|    |    └─LayerNorm: 3-16                   [-1, 256, 64]             128
|    |    └─Linear: 3-17                      [-1, 256, 128]            8,192
|    └─Downsample: 2-9                        [-1, 64, 8, 8]            --
|    |    └─Conv2d: 3-18                      [-1, 64, 8, 8]            36,928
|    └─Resnet_Block: 2-10                     [-1, 128, 8, 8]           --
|    |    └─Sequential: 3-19                  [-1, 128]                 16,512
|    |    └─Conv_block: 3-20                  [-1, 128, 8, 8]           74,112
|    |    └─Conv_block: 3-21                  [-1, 128, 8, 8]           147,840
|    |    └─Conv2d: 3-22                      [-1, 128, 8, 8]           8,320
|    └─UNET_Attention: 2-11                   [-1, 128, 8, 8]           --
|    |    └─LayerNorm: 3-23                   [-1, 64, 128]             256
|    |    └─MultiHeadAttention: 3-24          [-1, 64, 128]             65,536
|    |    └─LayerNorm: 3-25                   [-1, 64, 128]             256
|    |    └─Linear: 3-26                      [-1, 64, 256]             32,768
|    └─Conv2d: 2-12                           [-1, 128, 8, 8]           147,584
├─ModuleList: 1                               []                        --
|    └─Resnet_Block: 2-13                     [-1, 128, 8, 8]           --
|    |    └─Sequential: 3-27                  [-1, 128]                 16,512
|    |    └─Conv_block: 3-28                  [-1, 128, 8, 8]           147,840
|    |    └─Conv_block: 3-29                  [-1, 128, 8, 8]           147,840
|    |    └─Identity: 3-30                    [-1, 128, 8, 8]           --
|    └─UNET_Attention: 2-14                   [-1, 128, 8, 8]           --
|    |    └─LayerNorm: 3-31                   [-1, 64, 128]             256
|    |    └─MultiHeadAttention: 3-32          [-1, 64, 128]             65,536
|    |    └─LayerNorm: 3-33                   [-1, 64, 128]             256
|    |    └─Linear: 3-34                      [-1, 64, 256]             32,768
|    └─Resnet_Block: 2-15                     [-1, 64, 8, 8]            --
|    |    └─Sequential: 3-35                  [-1, 64]                  8,256
|    |    └─Conv_block: 3-36                  [-1, 64, 8, 8]            73,920
|    |    └─Conv_block: 3-37                  [-1, 64, 8, 8]            37,056
|    |    └─Conv2d: 3-38                      [-1, 64, 8, 8]            8,256
├─ModuleList: 1                               []                        --
|    └─ModuleList: 2                          []                        --
|    |    └─Resnet_Block: 3-39                [-1, 64, 8, 8]            168,448
|    |    └─UNET_Attention: 3-40              [-1, 64, 8, 8]            24,832
|    |    └─ConvTranspose2d: 3-41             [-1, 64, 8, 8]            36,928
|    └─ModuleList: 2                          []                        --
|    |    └─Resnet_Block: 3-42                [-1, 32, 8, 8]            54,528
|    |    └─UNET_Attention: 3-43              [-1, 32, 8, 8]            6,272
|    |    └─Upsample: 3-44                    [-1, 32, 16, 16]          16,416
|    └─ModuleList: 2                          []                        --
|    |    └─Resnet_Block: 3-45                [-1, 16, 16, 16]          14,720
|    |    └─UNET_Attention: 3-46              [-1, 16, 16, 16]          1,600
|    |    └─Upsample: 3-47                    [-1, 16, 32, 32]          4,112
├─GroupNorm: 1-5                              [-1, 16, 32, 32]          32
├─SiLU: 1-6                                   [-1, 16, 32, 32]          --
├─Conv2d: 1-7                                 [-1, 3, 32, 32]           435
===============================================================================================
Total params: 1,588,643
Trainable params: 1,588,643
Non-trainable params: 0
Total mult-adds (M): 32.79
===============================================================================================
Input size (MB): 0.04
Forward/backward pass size (MB): 5.51
Params size (MB): 6.06
Estimated Total Size (MB): 11.61
===============================================================================================

unet = to_device(unet, device)

"""
Function to Sample images from pure noise.

"""

def sampler(model, ddpm_sampler, xt: torch.Tensor, labels: torch.Tensor, cfg_scale: float, num_inference_steps: int, 
            save = False, train_noise_steps = 1000) -> torch.Tensor:
    
    """
    Args:
        model -> UNET model.
        ddpm_sampler -> DDPM.
        xt -> Random pure noise.
        labels -> Random labels.
        cfg_scale -> Weight to pay attention to the label signal for Classifier-free guidance approach. (7 >= cfg_scale <= 14)
        num_inference_steps -> Inference time steps.
        save -> Save generated images.
        train_noise_steps -> Train time steps.

    """
    
    index = 0
    model.eval()
    with torch.no_grad():
        
        step_ratio = train_noise_steps // num_inference_steps
        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].astype(np.int64)
        
        for i in timesteps:
                #t = (torch.ones(xt.size(0)) * i).unsqueeze(-1).long().to(xt.device)
                conditioned_predicted_noise = []
                t = torch.as_tensor(i).unsqueeze(0).to(xt.device)
                for img, label in zip(xt, labels):
                    pred_noise = model(img.unsqueeze(0), label, t)
                    conditioned_predicted_noise.append(pred_noise)
                predicted_noise = conditioned_predicted_noise = torch.stack(conditioned_predicted_noise, dim = 0).squeeze(1)
                # print(predicted_noise.size())
                if cfg_scale >= 7: 
                    unconditioned_predicted_noise = model(xt, None, t)
                    predicted_noise = cfg_scale * (conditioned_predicted_noise - unconditioned_predicted_noise) + unconditioned_predicted_noise
                #print(predicted_noise.size())

                xt = ddpm_sampler.sample(xt, predicted_noise, torch.as_tensor(i).to(xt.device))

                if save:
                    index += 1
                    show_and_save_images(xt, nmax = 64, t = i, index = index, save = save, show = False)
                     
            
    return xt, i

# Fixed Noise and Labels
xt = torch.randn((3, image_size[0], image_size[1], image_size[2]), generator = generator, device = device)
labels = torch.randint(0, 2, size = (3,), generator = generator, dtype = torch.long, device = device).flatten()
labels

tensor([0, 1, 0], device='cuda:0')

generated_image, i = sampler(unet, ddpm, xt, labels, cfg_scale = 7.2, num_inference_steps = num_inference_steps)

"""Generated Image before Training."""
show_and_save_images(generated_image, nmax = 64, t = 0, index = 0, save = True, show = True)

"""Train Function."""

def Train(model: nn.Module, noise_scheduler, optimizer, loss_criterion, n_timesteps: int, x: torch.Tensor, y: torch.Tensor):
    torch.cuda.empty_cache()
    t = torch.randint(0, n_timesteps, size = (x.size(0),), dtype = torch.long).to(x.device)
    
    # Diffusion process
    noise = torch.randn_like(x).to(x.device) # Epsilon
    noise_x = noise_scheduler.diffusion_process(x, t, noise)

    if np.random.random() < 0.2: # To uncondition the network based on the random probability
        y = None
        
    # Reverse process
    noise_pred = model(noise_x, y, t) # Epsilon_Theta

    loss = loss_criterion(noise_pred.reshape(-1, image_size[0]), noise.reshape(-1, image_size[0]))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss

"""Fit Function."""

def Fit(epochs: int, lr: float, n_timesteps: int, inference_steps: int):

    optimizer = torch.optim.Adam(unet.parameters(), lr = lr)
    loss_criterion = nn.MSELoss()
    Losses = []
    index = 0

    for epoch in range(epochs):
        unet.train()
        epoch_loss = []
        
        for x, y in tqdm(Train_Dl):
            loss = Train(unet, ddpm, optimizer, loss_criterion, n_timesteps, x, y)
            epoch_loss.append(loss)
        
        mean_loss = torch.stack(epoch_loss).mean().item()
        Losses.append(mean_loss)

        print("Epoch [{}], train_loss: {:.4f}".format(
               epoch + 1, mean_loss))
        
        #if epoch % 10 == 0:
        index += 1
        generated_images, time_step = sampler(unet, ddpm, xt, labels, cfg_scale = 7.2, num_inference_steps = inference_steps)
        show_and_save_images(generated_images, nmax = 64, t = time_step, index = index, save = True, show = False)
        print('Done Saving.....')
            
            
    return Losses

epochs = 200

"""Kick-start of Training."""
history = Fit(epochs, lr, num_training_steps, num_inference_steps)

  0%|          | 0/64 [00:00<?, ?it/s]

Epoch [1], train_loss: 0.6025
Done Saving.....

  0%|          | 0/64 [00:00<?, ?it/s]

Epoch [2], train_loss: 0.1983
Done Saving.....

  0%|          | 0/64 [00:00<?, ?it/s]

Epoch [3], train_loss: 0.1215
Done Saving.....

  0%|          | 0/64 [00:00<?, ?it/s]

Epoch [4], train_loss: 0.0936
Done Saving.....

  0%|          | 0/64 [00:00<?, ?it/s]

Epoch [5], train_loss: 0.0784
Done Saving.....

torch.save(unet.state_dict(), 'Diffusion.pth')

loss = history

x = range(len(loss))

plt.figure(figsize=(10,6))
plt.plot(x, loss ,'-')
plt.scatter(x, loss)
plt.xlabel('Epoch')
plt.ylabel('Diffusion model Loss')
plt.title('Diffusion model Loss over Time')
plt.show()

unet.load_state_dict(torch.load('Diffusion.pth'))

<All keys matched successfully>

# Random Noise and Labels
xt = torch.randn((batch_size, image_size[0], image_size[1], image_size[2]), device = device)
labels = torch.randint(0, 2, size = (batch_size,), dtype = torch.long, device = device).flatten()
labels

tensor([1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 1,
        1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0,
        1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0], device='cuda:0')

generated_images, last_t = sampler(unet, ddpm, xt, labels, cfg_scale = 7.5, num_inference_steps = num_inference_steps, save = True)

show_and_save_images(generated_images, nmax = 64, t = last_t, index = None, save = False, show = True)

# Random Noise and Female label
xt = torch.randn((batch_size, image_size[0], image_size[1], image_size[2]), generator = generator, device = device)
label = Label["female"]
labels = torch.ones(xt.size(0), dtype = torch.long, device = device).flatten() * label
labels

tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], device='cuda:0')

generated_images, _ = sampler(unet, ddpm, xt, labels, cfg_scale = 7.5, num_inference_steps = num_inference_steps, save = False)

show_and_save_images(generated_images, nmax = 64, t = 0, index = None, save = False, show = True)

# Random Noise and Male label
xt = torch.randn((batch_size, image_size[0], image_size[1], image_size[2]), generator = generator, device = device)
label = Label["male"]
labels = torch.ones(xt.size(0), dtype = torch.long, device = device).flatten() * label
labels

tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], device='cuda:0')

generated_images, _ = sampler(unet, ddpm, xt, labels, cfg_scale = 7.5, num_inference_steps = num_inference_steps, save = False)

show_and_save_images(generated_images, nmax = 64, t = 0, index = None, save = False, show = True)

# Random Noise and no labels
xt = torch.randn((batch_size, image_size[0], image_size[1], image_size[2]), generator = generator, device = device)
label = None
labels = [label] * batch_size
labels

[None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None]

generated_images, _ = sampler(unet, ddpm, xt, labels, cfg_scale = 0, num_inference_steps = num_inference_steps, save = False)

show_and_save_images(generated_images, nmax = 64, t = 0, index = None, save = False, show = True)

"""Create Video for Generated Images."""

vid_fname = 'ddpm_bitmoji_faces.mp4'

files = [os.path.join(sample_dir, f) for f in os.listdir(sample_dir) if 'generated-image' in f]
files.sort()

out = cv2.VideoWriter(vid_fname,cv2.VideoWriter_fourcc(*'MP4V'), 100, (640,663))
[out.write(cv2.imread(fname)) for fname in files]
out.release()

model_graph = draw_graph(unet, input_data = (img_dummy, labels_dummy, t), device = 'cpu', expand_nested = True)
model_graph.visual_graph

(process:28636): Pango-WARNING **: 16:17:42.979: couldn't load font "Linux libertine Not-Rotated 10", falling back to "Sans Not-Rotated 10", expect ugly output.

Denoising Diffusion Probabilistic Model [DDPM]¶

1. Dataset Creation and Visualization¶

2. Diffusion Scheduler¶

3. UNET¶

4. Train and Sampling Functions¶

5. Visualization of Loss¶

6. Inference¶

7. References¶